#include <cuda_runtime.h>
#include <device_launch_parameters.h>

#include "lastError.h"
#include "allocator.h"

namespace Image2D {

int Allocator::alloc2d(void* data[4], size_t linesize[4], int width, int height, ImageFormat format, int typesize)
{
    switch (format)
    {
    case ImageFormat::IMAGE_YUV420p:
        cudaMallocPitch(&data[0], &linesize[0], width * typesize,     height    );
        cudaMallocPitch(&data[1], &linesize[1], width * typesize / 2, height / 2);
        cudaMallocPitch(&data[2], &linesize[2], width * typesize / 2, height / 2);
        break;

    case ImageFormat::IMAGE_NV12:
        cudaMallocPitch(&data[0], &linesize[0], width * typesize, height    );
        cudaMallocPitch(&data[1], &linesize[1], width * typesize, height / 2);
        break;

    case ImageFormat::IMAGE_RGB:
    case ImageFormat::IMAGE_BGR:
        cudaMallocPitch(&data[0], &linesize[0], width * typesize * 3, height);
        // cudaMallocPitch(&data[1], &linesize[1], width * typesize, height);
        // cudaMallocPitch(&data[2], &linesize[2], width * typesize, height);
        break;

    case ImageFormat::IMAGE_RGBA:
    case ImageFormat::IMAGE_BGRA:
        cudaMallocPitch(&data[0], &linesize[0], width * typesize * 4, height);
        // cudaMallocPitch(&data[1], &linesize[1], width * typesize, height);
        // cudaMallocPitch(&data[2], &linesize[2], width * typesize, height);
        // cudaMallocPitch(&data[3], &linesize[3], width * typesize, height);
        break;

    default:
        cudaMallocPitch(&data[0], &linesize[0], width * typesize,     height    );
        break;
    }

    LastError::setLastError("error: alloc2d");
    return LastError::getLastErrorCode();
}

int Allocator::free2d(void* data[4])
{
    for (size_t i = 0; i < 4; i++)
    {
        if (data[i] != nullptr)
        {
            cudaFree(data[i]);
            data[i] = nullptr;
        }
    }

    LastError::setLastError("error: free2d");
   return LastError::getLastErrorCode();
}

int Allocator::copy2d(const void** srcData, const size_t* srcLinesize, void** dstData, size_t* dstLinesize,  int width, int height, ImageFormat format, int typesize)
{
    switch (format)
    {
    case ImageFormat::IMAGE_YUV420p:
        cudaMemcpy2D(dstData[0], dstLinesize[0], srcData[0], srcLinesize[0], width * typesize,     height,     cudaMemcpyDeviceToDevice);
        cudaMemcpy2D(dstData[1], dstLinesize[1], srcData[1], srcLinesize[1], width * typesize / 2, height / 2, cudaMemcpyDeviceToDevice);
        cudaMemcpy2D(dstData[2], dstLinesize[2], srcData[2], srcLinesize[2], width * typesize / 2, height / 2, cudaMemcpyDeviceToDevice);
        break;

    case ImageFormat::IMAGE_NV12:
        cudaMemcpy2D(dstData[0], dstLinesize[0], srcData[0], srcLinesize[0], width * typesize, height,     cudaMemcpyDeviceToDevice);
        cudaMemcpy2D(dstData[1], dstLinesize[1], srcData[1], srcLinesize[1], width * typesize, height / 2, cudaMemcpyDeviceToDevice);
        break;

    case ImageFormat::IMAGE_RGB:
    case ImageFormat::IMAGE_BGR:
        cudaMemcpy2D(dstData[0], dstLinesize[0], srcData[0], srcLinesize[0], width * typesize * 3, height, cudaMemcpyDeviceToDevice);
        // cudaMemcpy2D(dstData[1], dstLinesize[1], srcData[1], srcLinesize[1], width * typesize, height, cudaMemcpyDeviceToDevice);
        // cudaMemcpy2D(dstData[2], dstLinesize[2], srcData[2], srcLinesize[2], width * typesize, height, cudaMemcpyDeviceToDevice);
        break;

    case ImageFormat::IMAGE_RGBA:
    case ImageFormat::IMAGE_BGRA:
        cudaMemcpy2D(dstData[0], dstLinesize[0], srcData[0], srcLinesize[0], width * typesize * 4, height, cudaMemcpyDeviceToDevice);
        // cudaMemcpy2D(dstData[1], dstLinesize[1], srcData[1], srcLinesize[1], width * typesize, height, cudaMemcpyDeviceToDevice);
        // cudaMemcpy2D(dstData[2], dstLinesize[2], srcData[2], srcLinesize[2], width * typesize, height, cudaMemcpyDeviceToDevice);
        // cudaMemcpy2D(dstData[4], dstLinesize[4], srcData[4], srcLinesize[4], width * typesize, height, cudaMemcpyDeviceToDevice);
        break;

    default:
        cudaMemcpy2D(dstData[0], dstLinesize[0], srcData[0], srcLinesize[0], width * typesize, height, cudaMemcpyDeviceToDevice);
        break;
    }

    LastError::setLastError("error: copy2d");
    return LastError::getLastErrorCode();
}

}